#!/bin/sh

e2e_test_install() {
  NODE_LABEL_KEY="$(random_string)"
  NODE_LABEL_VALUE="$(random_string)"
  NODE_AFFINITY_VALUE="$(random_string)"
  NODE_AFFINITY_OPERATOR="$(random_string)"
  NODE_AFFINITY_KEY="$(random_string)"

  kubectl create namespace "$CLUSTER_NAMESPACE"

  create_or_replace_cluster "$CLUSTER_NAME" "$CLUSTER_NAMESPACE" 1 \
    --set-string "cluster.pods.scheduling.nodeSelector.$NODE_LABEL_KEY=$NODE_LABEL_VALUE" \
    --set-string "cluster.pods.scheduling.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].key=$NODE_AFFINITY_KEY" \
    --set-string "cluster.pods.scheduling.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].operator=$NODE_AFFINITY_OPERATOR" \
    --set-string "cluster.pods.scheduling.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].values[0]=$NODE_AFFINITY_VALUE" \
    --set "cluster.pods.scheduling.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].weight=1"\
    --set-string "cluster.pods.scheduling.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].preference.matchExpressions[0].key=$NODE_AFFINITY_KEY" \
    --set-string "cluster.pods.scheduling.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].preference.matchExpressions[0].operator=$NODE_AFFINITY_OPERATOR" \
    --set-string "cluster.pods.scheduling.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].preference.matchExpressions[0].values[0]=$NODE_AFFINITY_VALUE"
}

e2e_test_uninstall() {
  helm_cleanup_chart "$CLUSTER_NAME" "$CLUSTER_NAMESPACE"

  k8s_async_cleanup_namespace "$CLUSTER_NAMESPACE"
}

e2e_test() {
  run_test "cluster has specified node affinity required during scheduling" cluster_has_node_affinity_required_during_scheduling_properties
  run_test "cluster has specified node affinity preferred during scheduling" should_sgcluster_has_node_affinity_preferred_during_scheduling
  run_test "cluster's StatefulSet has specified required node affinity" check_sts_has_node_affinity_required_during_scheduling
  run_test "cluster's StatefulSet has specified preferred node affinity" check_sts_has_node_affinity_preferred_during_scheduling
}

cluster_has_node_affinity_required_during_scheduling_properties() {
  assert_cluster_node_affinity_required_during_scheduling_match_property "key" "$NODE_AFFINITY_KEY"
  assert_cluster_node_affinity_required_during_scheduling_match_property "operator" "$NODE_AFFINITY_OPERATOR"
  assert_cluster_node_affinity_required_during_scheduling_match_property "values[0]" "$NODE_AFFINITY_VALUE"
}

should_sgcluster_has_node_affinity_preferred_during_scheduling(){
  assert_cluster_node_affinity_preferred_during_scheduling_match_property "key" "$NODE_AFFINITY_KEY"
  assert_cluster_node_affinity_preferred_during_scheduling_match_property "operator" "$NODE_AFFINITY_OPERATOR"
  assert_cluster_node_affinity_preferred_during_scheduling_match_property "values[0]" "$NODE_AFFINITY_VALUE"

  if [ "$(kubectl get sgcluster -n "$CLUSTER_NAMESPACE" "$CLUSTER_NAME" -o json \
    | jq -r '.spec.pods.scheduling.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].weight')" \
    = "1" ]
  then
    success "cluster has the configured node affinity preferred during execution weight property"
  else
    fail "cluster has not the configured node affinity preferred during execution weight property"
  fi
}

check_sts_has_node_affinity_required_during_scheduling() {
  local RESTUL EXIT_CODE
  echo "Waiting for sts with node affinity $CLUSTER_NAMESPACE $CLUSTER_NAME"
  try_function wait_until eval "$(cat << 'EOF'
[ "$(kubectl get sts -n "$CLUSTER_NAMESPACE" "$CLUSTER_NAME" -o json \
  | jq -r ".spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].key")" \
  = "$NODE_AFFINITY_KEY" ]
EOF
    )"
  if "$RESULT"
  then
    success "cluster's StatefulSet has the configured node affinity required during execution key match"
  else
    fail "cluster's StatefulSet has not the configured node affinity required during execution key match"
  fi

  assert_sts_node_affinity_required_during_scheduling_match_property "operator" "$NODE_AFFINITY_OPERATOR"
  assert_sts_node_affinity_required_during_scheduling_match_property "values[0]" "$NODE_AFFINITY_VALUE"
}

check_sts_has_node_affinity_preferred_during_scheduling(){
    local RESTUL EXIT_CODE
  echo "Waiting for sts with preferred node affinity $CLUSTER_NAMESPACE $CLUSTER_NAME"
  try_function wait_until eval "$(cat << 'EOF'
[ "$(kubectl get sts -n "$CLUSTER_NAMESPACE" "$CLUSTER_NAME" -o json \
  | jq -r ".spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].key")" \
  = "$NODE_AFFINITY_KEY" ]
EOF
    )"
  if "$RESULT"
  then
    success "cluster's StatefulSet has the configured node affinity preferred during execution key match"
  else
    fail "cluster's StatefulSet has not the configured node affinity preferred during execution key match"
  fi

  assert_sts_node_affinity_preferred_during_scheduling_match_property "operator" "$NODE_AFFINITY_OPERATOR"
  assert_sts_node_affinity_preferred_during_scheduling_match_property "values[0]" "$NODE_AFFINITY_VALUE"

  if [ "$(kubectl get sts -n "$CLUSTER_NAMESPACE" "$CLUSTER_NAME" -o json \
      | jq -r ".spec.template.spec.affinity.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].weight")" \
    = "1" ]
  then
    success "cluster's StatefulSet has the configured node affinity preferred during execution weight property"
  else
    fail "cluster's StatefulSet has not the configured node affinity preferred during execution weight property"
  fi
}

assert_sts_node_affinity_preferred_during_scheduling_match_property(){
  local PROPERTY="$1"
  local PROPERTY_VALUE="$2"

  if [ "$(kubectl get sts -n "$CLUSTER_NAMESPACE" "$CLUSTER_NAME" -o json \
      | jq -r ".spec.template.spec.affinity.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].preference.matchExpressions[0].$PROPERTY")" \
    = "$PROPERTY_VALUE" ]
  then
    success "StatefulSet has the configured node affinity preferred during scheduling with match for $PROPERTY set to $PROPERTY_VALUE"
  else
    fail "StatefulSet has not the configured node affinity preferred during scheduling with match for $PROPERTY set to $PROPERTY_VALUE"
  fi
}
assert_sts_node_affinity_required_during_scheduling_match_property(){
  local PROPERTY="$1"
  local PROPERTY_VALUE="$2"

  if [ "$(kubectl get sts -n "$CLUSTER_NAMESPACE" "$CLUSTER_NAME" -o json \
      | jq -r ".spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].$PROPERTY")" \
    = "$PROPERTY_VALUE" ]
  then
    success "StatefulSet has the configured node affinity required during scheduling with match for $PROPERTY set to $PROPERTY_VALUE"
  else
    fail "StatefulSet has not the configured node affinity required during scheduling with match for $PROPERTY set to $PROPERTY_VALUE"
  fi
}

assert_cluster_node_affinity_preferred_during_scheduling_match_property(){
  local PROPERTY="$1"
  local PROPERTY_VALUE="$2"

  if [ "$(kubectl get sgcluster -n "$CLUSTER_NAMESPACE" "$CLUSTER_NAME" -o json \
      | jq -r ".spec.pods.scheduling.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution[0].preference.matchExpressions[0].$PROPERTY")" \
    = "$PROPERTY_VALUE" ]
  then
    success "cluster has the configured node affinity preferred during scheduling with match for $PROPERTY set to $PROPERTY_VALUE"
  else
    fail "cluster has not the configured node affinity preferred during scheduling with match for $PROPERTY set to $PROPERTY_VALUE"
  fi
}
assert_cluster_node_affinity_required_during_scheduling_match_property(){
  local PROPERTY="$1"
  local PROPERTY_VALUE="$2"

  if [ "$(kubectl get sgcluster -n "$CLUSTER_NAMESPACE" "$CLUSTER_NAME" -o json \
      | jq -r ".spec.pods.scheduling.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].$PROPERTY")" \
    = "$PROPERTY_VALUE" ]
  then
    success "cluster has the configured node affinity required during scheduling with match for $PROPERTY set to $PROPERTY_VALUE"
  else
    fail "cluster has not the configured node affinity required during scheduling with match for $PROPERTY set to $PROPERTY_VALUE"
  fi
}
